import time
import jieba_fast as jieba
from tqdm import tqdm

ARTICLE_FILE = "datas/新闻标题数据集/train_text.txt"
SUMMARRY_FILE = "datas/新闻标题数据集/train_label.txt"

TRAIN_FILE = "./datas/train_art_summ_prep.txt"
VAL_FILE = "./datas/val_art_summ_prep.txt"

"""
第一部分是对原始数据进行分词，划分训练集测试集，并保存文件。
"""


def timer(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        r = func(*args, **kwargs)
        end = time.time()
        cost = end - start
        print(f"Cost time: {cost} s")
        return r

    return wrapper


@timer
def load_data(filename):
    """加载数据文件，对文本进行分词"""
    data_list = []
    with open(filename, "r", encoding="utf-8") as f:
        for line in tqdm(f, desc="loading datas"):
            # jieba.enable_parallel()
            words = jieba.cut(line.strip())
            word_list = list(words)
            # jieba.disable_parallel()
            data_list.append(" ".join(word_list).strip())
    return data_list


def build_train_val(article_data, summary_data, train_num=600_000):
    """划分训练和验证数据"""
    train_list = []
    val_list = []
    n = 0
    for text, summ in zip(article_data, summary_data):
        n += 1
        if n <= train_num:
            train_list.append(text)
            train_list.append(summ)
        else:
            val_list.append(text)
            val_list.append(summ)
    return train_list, val_list


def save_file(filename, li):
    """预处理后的数据保存到文件"""
    with open(filename, "w+", encoding="utf-8") as f:
        for item in li:
            f.write(item + "\n")
    print(f"Save {filename} ok.")


if __name__ == "__main__":
    article_data = load_data(ARTICLE_FILE)  # 大概耗时10分钟
    summary_data = load_data(SUMMARRY_FILE)
    TRAIN_SPLIT = 600_000
    train_list, val_list = build_train_val(
        article_data, summary_data, train_num=TRAIN_SPLIT
    )
    save_file(TRAIN_FILE, train_list)
    save_file(VAL_FILE, val_list)
